# -*- coding: utf-8 -*-
"""
水库调度规程智能问答系统 - Streamlit前端界面（真流式输出版本）
支持实时流式输出和对话历史
"""

import sys
import os
from pathlib import Path
import streamlit as st
from typing import Generator
import time

# 将项目根目录添加到路径
project_root = Path(__file__).resolve().parent
sys.path.insert(0, str(project_root))

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate

from D_db.core.db_manager import DBManager
from D_db.core.tokenizer import Tokenizer
from D_db.tools.prompt_utils import build_llm_context_from_qa

# 加载环境变量
load_dotenv()

# 阿里云百炼配置
DASHSCOPE_CONFIG = {
    "api_key": os.getenv('DASHSCOPE_API_KEY'),
    "base_url": os.getenv('BASE_URL'),
    "model": "qwen-long"
}

# 页面配置
st.set_page_config(
    page_title="水库调度规程智能问答系统",
    page_icon="💧",
    layout="wide",
    initial_sidebar_state="expanded"
)

# 自定义CSS样式
st.markdown("""
<style>
    .main-header {
        font-size: 2.5rem;
        color: #1E88E5;
        text-align: center;
        padding: 1rem 0;
        border-bottom: 3px solid #1E88E5;
        margin-bottom: 2rem;
    }
    .stChatMessage {
        padding: 1rem;
        border-radius: 10px;
        margin-bottom: 1rem;
    }
    .info-box {
        background-color: #FFF3E0;
        padding: 1rem;
        border-radius: 8px;
        border-left: 4px solid #FF9800;
        margin: 1rem 0;
    }
</style>
""", unsafe_allow_html=True)

class StreamingQueryTester:
    """支持流式输出的查询测试器"""
    
    def __init__(self):
        """初始化"""
        self.db_manager = DBManager()
        self.tokenizer = Tokenizer()
        
        # 初始化阿里云百炼模型
        self.llm = ChatOpenAI(
            api_key=DASHSCOPE_CONFIG["api_key"],
            base_url=DASHSCOPE_CONFIG["base_url"],
            model=DASHSCOPE_CONFIG["model"],
            temperature=0.7,
            streaming=True  # 启用流式输出
        )
        
        # 创建提示模板
        self.answer_template = PromptTemplate(
            input_variables=["question", "context"],
            template="""你是一个专业的水库调度规程问答助手。请根据以下参考信息，回答用户的问题。

用户问题：
{question}

参考信息：
{context}

要求：
1. 基于参考信息提供准确、完整的回答
2. 如果参考信息不足以回答问题，请明确指出
3. 回答要简洁明了，直接针对问题给出答案
4. 不要在回答中提及"根据参考信息"等字样
5. 使用专业、客观的语气

请提供你的回答：
"""
        )
        
        self.answer_chain = self.answer_template | self.llm
    
    def query_database(self, test_question: str, limit: int = 20):
        """查询数据库"""
        try:
            self.db_manager.connect()
            
            # 分词
            token_data = self.tokenizer.tokenize_question(test_question)
            tokens = token_data.get('tokens', [])
            category_id = token_data.get('category_id')
            
            if not tokens:
                return None, []
            
            # 查询
            results = self.db_manager.query_qa(tokens, limit, category_id)
            
            # 构建上下文
            context_for_llm, qa_results = build_llm_context_from_qa(results)
            
            return context_for_llm, qa_results
            
        finally:
            self.db_manager.close()
    
    def stream_answer(self, question: str, context: str):
        """流式生成回答"""
        try:
            # 使用 stream 方法流式输出
            for chunk in self.answer_chain.stream({"question": question, "context": context}):
                if hasattr(chunk, 'content'):
                    yield chunk.content
                else:
                    yield str(chunk)
        except Exception as e:
            yield f"生成回答时出错: {str(e)}"

def initialize_session_state():
    """初始化 session state"""
    if "messages" not in st.session_state:
        st.session_state.messages = []
    
    if "tester" not in st.session_state:
        try:
            st.session_state.tester = StreamingQueryTester()
            st.session_state.tester_initialized = True
        except Exception as e:
            st.session_state.tester = None
            st.session_state.tester_initialized = False
            st.session_state.init_error = str(e)

def format_qa_results(qa_results: list) -> str:
    """格式化查询结果"""
    if not qa_results:
        return ""
    
    result_text = "\n\n### 📚 参考的知识库内容\n\n"
    
    for i, item in enumerate(qa_results[:5], 1):
        result_text += f"**{i}. {item.get('question', '')}**\n\n"
        result_text += f"*{item.get('answer', '')}*\n\n"
        
        if item.get('category_id'):
            result_text += f"🏷️ 水库ID: `{item['category_id']}`  "
        if item.get('match_count'):
            result_text += f"| 🎯 匹配度: `{item['match_count']}`\n\n"
        
        result_text += "---\n\n"
    
    return result_text

def main():
    """主函数"""
    # 初始化
    initialize_session_state()
    
    # 标题
    st.markdown('<h1 class="main-header">💧 水库调度规程智能问答系统</h1>', unsafe_allow_html=True)
    
    # 侧边栏
    with st.sidebar:
        st.header("⚙️ 系统信息")
        
        if st.session_state.tester_initialized:
            st.success("✅ 系统已就绪（流式输出模式）")
            
            st.info("""
            **支持的水库：**
            - 大伙房水库
            - 观音阁水库
            - 葠窝水库
            - 柴河水库
            - 汤河水库
            - 白石水库
            - 清河水库
            - 石佛寺水库
            - 闹德海水库
            """)
            
            st.divider()
            
            # 查询设置
            st.subheader("🔧 查询设置")
            limit = st.slider("返回结果数量", min_value=5, max_value=50, value=20, step=5)
            st.session_state.query_limit = limit
            
            st.divider()
            
            # 示例问题
            st.subheader("💡 示例问题")
            example_questions = [
                "大伙房水库的总库容是多少？",
                "大伙房水库的正常高水位是多少？",
                "大伙房水库的防洪标准是什么？",
                "大伙房水库的主要功能有哪些？",
                "大伙房水库的流域面积是多少？"
            ]
            
            for question in example_questions:
                if st.button(question, key=f"example_{question}", use_container_width=True):
                    st.session_state.selected_question = question
            
            st.divider()
            
            # 清空对话
            if st.button("🗑️ 清空对话历史", use_container_width=True):
                st.session_state.messages = []
                st.rerun()
        else:
            st.error("❌ 系统初始化失败")
            st.error(f"错误信息：{st.session_state.get('init_error', '未知错误')}")
            
            st.markdown("""
            <div class="info-box">
            <strong>⚠️ 请检查：</strong><br>
            1. 数据库是否已创建和导入数据<br>
            2. .env 文件是否配置正确<br>
            3. DASHSCOPE_API_KEY 是否有效<br>
            4. PostgreSQL 服务是否启动
            </div>
            """, unsafe_allow_html=True)
            
            if st.button("🔄 重新初始化", use_container_width=True):
                for key in ["tester", "tester_initialized", "init_error"]:
                    if key in st.session_state:
                        del st.session_state[key]
                st.rerun()
    
    # 主界面
    if not st.session_state.tester_initialized:
        st.warning("⚠️ 系统未就绪，请查看左侧边栏的错误信息")
        return
    
    # 处理选中的示例问题
    if hasattr(st.session_state, 'selected_question'):
        prompt = st.session_state.selected_question
        del st.session_state.selected_question
        
        # 添加到历史并处理
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        # 查询并生成回答
        try:
            limit = st.session_state.get('query_limit', 20)
            context, qa_results = st.session_state.tester.query_database(prompt, limit=limit)
            
            if not context:
                response = "抱歉，没有找到相关信息。请尝试重新表述您的问题。"
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": response,
                    "qa_results": []
                })
            else:
                # 生成完整回答（不流式）
                full_response = ""
                for chunk in st.session_state.tester.stream_answer(prompt, context):
                    full_response += chunk
                
                # 添加到历史
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": full_response,
                    "qa_results": qa_results
                })
        except Exception as e:
            error_msg = f"❌ 查询出错：{str(e)}"
            st.session_state.messages.append({
                "role": "assistant",
                "content": error_msg,
                "qa_results": []
            })
        
        # 重新运行以显示历史记录
        st.rerun()
    
    # 显示对话历史
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
            
            # 显示参考资料
            if message["role"] == "assistant" and message.get("qa_results"):
                with st.expander("🔍 查看参考的知识库内容", expanded=False):
                    formatted_results = format_qa_results(message["qa_results"])
                    st.markdown(formatted_results)
    
    # 输入框
    prompt = st.chat_input("请输入您的问题...")
    
    # 处理用户输入
    if prompt:
        # 添加用户消息到历史
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        # 查询并生成回答
        try:
            # 查询数据库
            limit = st.session_state.get('query_limit', 20)
            context, qa_results = st.session_state.tester.query_database(prompt, limit=limit)
            
            if not context:
                response = "抱歉，没有找到相关信息。请尝试重新表述您的问题。"
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": response,
                    "qa_results": []
                })
            else:
                # 流式生成回答
                full_response = ""
                for chunk in st.session_state.tester.stream_answer(prompt, context):
                    full_response += chunk
                
                # 添加到历史
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": full_response,
                    "qa_results": qa_results
                })
        
        except Exception as e:
            error_msg = f"❌ 查询出错：{str(e)}"
            st.session_state.messages.append({
                "role": "assistant",
                "content": error_msg,
                "qa_results": []
            })
        
        # 重新运行以显示历史记录
        st.rerun()

if __name__ == "__main__":
    main()

